# -*- coding: utf-8 -*-
"""
Created on Thu Jun 21 16:25:23 2012

@author: nzhao
"""

from traits.api import HasTraits, Instance, Enum, String, Bool, Float, Array, Button, Int, List
from traitsui.api import View, Group, HGroup, Item, RangeEditor, ArrayEditor, InstanceEditor

from enable.api import Component, ComponentEditor
from chaco.api import VPlotContainer, HPlotContainer

import numpy as np

from phy.control.CPMG import CPMG
from math_qm.CoordianteSystem import CoordinateSystem
from math_qm.vector import Cartesian_Vector, Sphere_Vector
from misc.multilines import OverlappingPlotContainer

control_group = HGroup(
                    Group(
                        Item('axes_b', editor = InstanceEditor( name = 'axes_list', editable = False), label='axes', style = 'custom'),
                        Item('B_vector', style = 'custom', show_label = False),
                        show_border = True, label='B field'
                    ),
                    Group(
                        HGroup(
                               Item('time_list_str',width = 300, label = 'Time points:'),
                               Item('transition', style = 'custom'),
                               ),
                        HGroup(
                               Item('nPulse', editor = RangeEditor(low=1, high_name = 'maxNpulse', mode = 'spinner'),width = 100),
                               Item('maxNpulse'),
                               ),
                        show_border = True, label='Control', springy = True
                    )
                )
auto_update_group = HGroup(
                           Item('is_auto_update', label = 'Auto Update'),
                           Item('do_cpmg_cal', show_label=False, enabled_when = 'not auto_update'),
                    )
cpmg_group = Group(
                auto_update_group,
                control_group,
                HGroup(
                        Group(
                            Item('A_vector', style='custom', show_label=False),
                        ),
                        Group(
                            Item('_'),
                            Item('fieldPair', style='readonly', editor=ArrayEditor(width=150, format_str='%8.3f')),
                        ),
                        show_border=True, label='hyperfine'
                    ),
                Item('plot123', editor=ComponentEditor(), show_label=False),
                label='CPMG-Single sin-1/2'
            )

class cpmg_spin_half_hf(HasTraits):
    is_auto_update   = Bool(False)
    do_cpmg_cal      = Button(label='Calculate CPMG')

    cpmg_calculator  = Instance(CPMG)

    axes_b      = Instance( CoordinateSystem )
    axes_list        = List( CoordinateSystem )
    
    B_vector = Instance(Sphere_Vector, ())
    A_vector = Instance(Cartesian_Vector, ())
    
    fieldPair = Array
    
    s1         = Float(0.0)
    s2         = Float(1.0)
    transition = Enum('SQT', 'DQT')
    
    time_list_str = String('np.linspace(0.,10, 301)', auto_set=False, enter_set = False)
    nPulse = Int(1)
    maxNpulse=Int(100)
    
    plot1   = Instance(Component)
    plot2   = Instance(Component)
    plot3   = Instance(Component)    
    plot23  = Instance(HPlotContainer)
    plot123 = Instance(VPlotContainer)

    def __init__(self, axes = None):
        self.A_vector = Cartesian_Vector(name='HF Field')
        self.B_vector = Sphere_Vector(name='B Field')
        self.update_fields()        
            
        if axes is not None:
            self.axes_list = axes
            self.axes_b = axes[0]
        else:
            self.axes_b = CoordinateSystem()
        
        self.cpmg_calculator = CPMG(spin_num = 1)
       
        self.A_vector.on_trait_change(self.update_fields, 'x.value, y.value, z.value')
        self.B_vector.on_trait_change(self.update_fields, 'x.value, y.value, z.value')


        self.plot1=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)
        self.plot2=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)
        self.plot3=OverlappingPlotContainer(padding=30, fill_padding=True, bgcolor="white", use_backbuffer=True)

    def _axes_b_changed(self):
        self.update_fields()
        
    def _transition_changed(self):
        if self.transition == 'SQT':
            self.s1, self.s2 = 0.0, 1.0
        else:
            self.s1, self.s2 = -1.0, 1.0

        self.update_fields()
    
    def update_fields(self):
        bx, by, bz = self.B_vector.x.value, self.B_vector.y.value, self.B_vector.z.value
        Ax, Ay, Az = self.A_vector.x.value, self.A_vector.y.value, self.A_vector.z.value
        
        b_vector  = np.array([bx, by, bz]) 
        hf_vector = np.array([Ax, Ay, Az]) 
                
        h0 = b_vector + self.s1 * hf_vector
        h1 = b_vector + self.s2 * hf_vector
        
        self.fieldPair =[h0, h1]
        if self.is_auto_update:
            self.plot_cpmg()

            
    def _do_cpmg_cal_fired(self):
            self.plot_cpmg()
    def _time_list_str_changed(self):   
        if self.is_auto_update:
            self.plot_cpmg()
    def _nPulse_changed(self):
        if self.is_auto_update:
            self.plot_cpmg()

    def calculate_data(self, xlist):
        ylist = self.cpmg_calculator.coherence(timelist = xlist, pseudo_field= self.fieldPair, pulse_num= self.nPulse)
        return ylist

    def plot_cpmg(self):
        xlist = eval(self.time_list_str)
        ylist = self.calculate_data(xlist)
        
        self.plot1.clear_data()
        self.plot2.clear_data()
        self.plot3.clear_data()

        for i in range(4):
            self.plot1.add_data([xlist, ylist[:, 0, i]])
            self.plot2.add_data([xlist, ylist[:, 1, i]])
            self.plot3.add_data([xlist, ylist[:, 2, i]])
        self.plot1.plot_data(['black', 'red', 'green', 'blue'])
        self.plot2.plot_data(['black', 'red', 'green', 'blue'])
        self.plot3.plot_data(['black', 'red', 'green', 'blue'])
        
        self.plot23  = HPlotContainer(self.plot2, self.plot3, stack_order="left_to_right")
        self.plot123 = VPlotContainer(self.plot1, self.plot23, stack_order="top_to_bottom")
        self.plot2.components[0].index =self.plot1.components[0].index
        self.plot3.components[0].index =self.plot1.components[0].index

    cpmg_view = View(
        cpmg_group,
        title     = 'CPMG',
        width     = .8,
        height    = .9,
        resizable = True,
        buttons   = [ 'OK' ],
        kind      = 'live'
    )
     
if __name__ == '__main__':
    cpmg1 = cpmg_spin_half_hf()
    cpmg1.configure_traits()
